FROM nvidia/cuda:12.9.0-cudnn-runtime-ubuntu24.04

# 修复链接库的问题
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH
ENV PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
ENV CUDA_HOME=/usr/local/cuda

RUN cp /etc/apt/sources.list.d/ubuntu.sources /etc/apt/sources.list.d/ubuntu.sources.bak && \
    sed -i 's/archive.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list.d/ubuntu.sources && \
    export DEBIAN_FRONTEND=noninteractive && \
    apt-get update -y && \
    apt-get install wget build-essential python3-pip python3-venv -y

# 安装系统图形库（解决 OpenCV/PyMuPDF 依赖）
RUN export DEBIAN_FRONTEND=noninteractive  && \
    apt-get install -y \
    libgomp1 \
    libgl1-mesa-* \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6

# 解决 Cannot load cudnn shared library. Cannot invoke method cudnnGetVersion 的错误
RUN ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so.9 /usr/lib/x86_64-linux-gnu/libcudnn.so && \
    ln -s /usr/local/cuda/targets/x86_64-linux/lib/libcublas.so.12 /usr/lib/x86_64-linux-gnu/libcublas.so && \
    ln -s /usr/lib/x86_64-linux-gnu/libnccl.so.2 /usr/lib/x86_64-linux-gnu/libnccl.so

# 清理
RUN apt-get clean && \
    rm -rf /var/lib/apt/lists/* \
    rm -rf /tmp/*

CMD ["/bin/sh"]